# -*- coding: utf-8 -*-
"""
Created on Wed Nov 22 10:43:57 2017

@author: xuanlei
"""

from get_captcha import Vocab
import tf_cnn as tfcnn
import tensorflow as tf
import numpy as np






#==============================================================================
# Hyper-Para Setting
#==============================================================================


input_size = 60*160
height = 60
width = 160
len_char = 4
len_char_set = 10+26+26
LR = 0.001
c1_size = 32
c2_size = 64
c3_size = 64
h1_size = 1024
output_size = len_char_set*len_char
batch_size = 64


# =============================================================================
# get data
# =============================================================================


def convert2gray(img):
	if len(img.shape) > 2:
		gray = np.mean(img, -1)
		# 上面的转法较快，正规转法如下
		# r, g, b = img[:,:,0], img[:,:,1], img[:,:,2]
		# gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
		return gray
	else:
		return img




def get_data_batch(size=128):
    voc = Vocab()
    batch_x = np.zeros([size, height*width])
    batch_y = np.zeros([size, output_size])
    def get_data():
        while 1:
            text,data,_ = voc.gen_captcha()
            if data.shape == (60, 160, 3):
                return text,data
            else:
                pass
            
    for i in  range(size):
        text,data = get_data()
        gray_data = convert2gray(data)
        batch_x[i,:] = gray_data.flatten() / 255 # (image.flatten()-128)/128  mean为0
        batch_y[i,:] = voc.text_to_one_hot(text)
    return batch_x, batch_y
            



    
def train_cnn():
    print('..............>>开始构建模型')
    tf.reset_default_graph()
    model = tfcnn.CNN(input_size,output_size,height,width,len_char,len_char_set,c1_size, c2_size, c3_size,h1_size,LR,batch_size)
    saver = tf.train.Saver()
    print('..............>>模型构建完毕，开始训练')
    with tf.Session() as sess:
        init_op = tf.global_variables_initializer()
        sess.run(init_op)
        merged = tf.summary.merge_all()
        writer = tf.summary.FileWriter("cnn_map/", sess.graph)
        with tf.variable_scope('scope', reuse=True):
            max_epoch = 15000
            epoch = 0
            while epoch < max_epoch:
                batch_x, batch_y = get_data_batch(batch_size)
                feed_dict_train = {model.xs: batch_x, model.ys: batch_y,model.keep_prob: np.array(0.75,  dtype='float32'), model.train_phase: True}
                _, cost, train_acc = sess.run([model.train_op, model.cost, model.accuracy], feed_dict=feed_dict_train)
                rs = sess.run(merged,feed_dict=feed_dict_train)
                writer.add_summary(rs, epoch)
                print('train_acc:  %f'%train_acc)
                if epoch % 10 == 0:
                    batch_x_test, batch_y_test = get_data_batch(64)
                    test_acc = sess.run(model.accuracy, feed_dict={model.xs: batch_x_test, model.ys: batch_y_test, model.keep_prob:np.array(1, dtype='float32'),model.train_phase: False})
                    print('...............>>>now give test result !')
                    print(epoch,test_acc)
                    if test_acc > 0.70 :
                        saver.save(sess, "crack_capcha.model", global_step=epoch)
                        break
                epoch += 1

                

        
if __name__ == '__main__':
    train_cnn()